import pandas as pd
import os
from itertools import combinations


def read_and_clean_csv(file_path, unique_columns):
    """
    读取CSV文件，并自动清理指定列上的重复数据。

    Args:
        file_path (str): CSV文件的路径。
        unique_columns (list): 组合起来应具有唯一性的列名列表。

    Returns:
        DataFrame: 清理掉重复行后的DataFrame。
    """
    try:
        df = pd.read_csv(file_path)

        # 检查文件是否为空
        if df.empty:
            print(f"警告: 文件 {os.path.basename(file_path)} 是空的，已跳过。")
            return None

        # 检查文件是否包含所有必需的列
        if not set(unique_columns).issubset(df.columns):
            missing_cols = set(unique_columns) - set(df.columns)
            print(f"错误: 文件 {os.path.basename(file_path)} 缺少必需的列: {missing_cols}。已跳过。")
            return None

        # --- 核心修改：从检查重复变为自动去重 ---
        # 计算去重前的行数
        original_count = len(df)
        # 按指定列去重，保留第一次出现的行
        df.drop_duplicates(subset=unique_columns, keep='first', inplace=True)
        # 计算去重后的行数
        cleaned_count = len(df)

        # 如果有重复行被删除，打印提示信息
        if original_count > cleaned_count:
            removed_count = original_count - cleaned_count
            print(f"提示: 文件 {os.path.basename(file_path)} 中发现 {removed_count} 行重复数据，已自动清理。")

        # 将唯一列设置为索引，方便后续比较
        df.set_index(unique_columns, inplace=True)
        return df

    except FileNotFoundError:
        print(f"错误: 文件 {file_path} 未找到。")
        return None
    except Exception as e:
        print(f"读取文件 {os.path.basename(file_path)} 时发生未知错误: {e}")
        return None


def compare_and_save_diffs(df1, name1, df2, name2, output_dir):
    """
    比较两个DataFrame并将差异保存到CSV文件。
    (此函数与之前版本相同，无需修改)
    """
    if df1 is None or df2 is None:
        print(f"比较 {os.path.basename(name1)} 和 {os.path.basename(name2)} 时，其中一个或两个文件无效，已跳过。")
        return

    # 使用外连接找到所有在一个DataFrame中存在但在另一个中不存在的行
    merged = df1.merge(df2, how='outer', left_index=True, right_index=True, indicator=True)

    # 筛选出仅存在于第一个文件的行
    only_in_df1 = merged[merged['_merge'] == 'left_only'].drop(columns=['_merge'])
    # 筛选出仅存在于第二个文件的行
    only_in_df2 = merged[merged['_merge'] == 'right_only'].drop(columns=['_merge'])

    # 如果没有差异，则不创建文件
    if only_in_df1.empty and only_in_df2.empty:
        print(f"文件 {os.path.basename(name1)} 和 {os.path.basename(name2)} 完全相同，没有差异。")
        return

    # 准备输出DataFrame，添加一个列来指示差异来源
    diffs_df1 = only_in_df1.reset_index()
    diffs_df1['差异位置'] = os.path.basename(name1)

    diffs_df2 = only_in_df2.reset_index()
    diffs_df2['差异位置'] = os.path.basename(name2)

    # 合并两个差异DataFrame
    all_diffs = pd.concat([diffs_df1, diffs_df2], ignore_index=True)

    # 构造输出文件名
    output_filename = f"{os.path.basename(name1)}_vs_{os.path.basename(name2)}_diff.csv"
    output_path = os.path.join(output_dir, output_filename)

    # 保存到CSV文件
    all_diffs.to_csv(output_path, index=False, encoding='utf-8-sig')
    print(f"已将 {os.path.basename(name1)} 和 {os.path.basename(name2)} 的差异保存到: {output_path}")


def main():
    # --- 配置 ---
    input_dir = 'flaws'
    output_dir = 'comparison_results'
    unique_columns = ['pro_line_id', 'product_id', 'device_id', 'area_id']

    if not os.path.isdir(input_dir):
        print(f"错误: 输入文件夹 '{input_dir}' 不存在。请在脚本同级目录下创建该文件夹并放入CSV文件。")
        return

    os.makedirs(output_dir, exist_ok=True)

    csv_files = [os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.lower().endswith('.csv')]

    if not csv_files:
        print(f"在 '{input_dir}' 文件夹中没有找到任何CSV文件。")
        return

    print(f"找到 {len(csv_files)} 个CSV文件，准备进行两两对比...\n")

    # 读取并**清理**所有CSV文件
    dataframes = {}
    for file in csv_files:
        # 调用修改后的函数
        df = read_and_clean_csv(file, unique_columns)
        if df is not None:
            dataframes[file] = df

    if not dataframes:
        print("\n没有有效的CSV文件可供比较。程序已退出。")
        return

    # 两两对比
    for file1, file2 in combinations(dataframes.keys(), 2):
        compare_and_save_diffs(dataframes[file1], file1, dataframes[file2], file2, output_dir)

    print("\n所有文件对比完成！")


if __name__ == '__main__':
    main()